import chromadb
import os
from openai import OpenAI
from rag.knowledge_build import text_embeddings
from tools.myes import get_es_client, search_es_for_advice


def query_chroma(query_text, collection_name="employee_manual", n_results=5):
    """
    根据用户问题在ChromaDB中检索最相似的文本块

    参数:
        query_text (str): 用户的问题文本
        collection_name (str): Chroma集合名称
        n_results (int): 返回结果数量
    返回:
        dict: 检索结果
    """
    # 初始化ChromaDB客户端（连接到持久化存储路径）
    persist_directory = "./chroma_db"
    chroma_client = chromadb.PersistentClient(path=persist_directory)
    
    # 获取指定集合（不传入嵌入函数，避免与已存在的配置冲突）
    collection = chroma_client.get_collection(name=collection_name)

    # 为查询文本生成嵌入向量
    query_embeddings = text_embeddings([query_text])
    
    # 在集合中查询与问题最相似的文本块
    results = collection.query(
        query_embeddings=query_embeddings,  # 使用预生成的查询向量
        n_results=n_results,  # 返回最相似的结果数量
        include=['embeddings', 'documents', 'metadatas', 'distances']  # 指定返回的字段
    )
    
    return results


# def query_elasticsearch(query_text, index_name="health_advice", n_results=5):
#     """
#     根据用户问题在Elasticsearch中检索相关文本块
#
#     参数:
#         query_text (str): 用户的问题文本
#         index_name (str): Elasticsearch索引名称
#         n_results (int): 返回结果数量
#     返回:
#         dict: 检索结果
#     """
#     # 初始化Elasticsearch客户端
#     es_client = get_es_client()
#
#     # 在Elasticsearch中搜索
#     response = search_es_for_advice(es_client, query_text, index_name=index_name)
#
#     # 提取前n_results个结果
#     hits = response['hits']['hits'][:n_results]
#
#     # 构造与Chroma相似的结果格式
#     results = {
#         'ids': [hit['_id'] for hit in hits],
#         'documents': [hit['_source']['text_content'] for hit in hits],
#         'scores': [hit['_score'] for hit in hits]
#     }
#
#     return results
#
#
# def reciprocal_rank_fusion(chroma_results, es_results, k=60):
#     """
#     使用Reciprocal Rank Fusion (RRF) 算法融合两个检索系统的分数
#
#     参数:
#         chroma_results (dict): Chroma检索结果
#         es_results (dict): Elasticsearch检索结果
#         k (int): RRF参数，通常设为60
#
#     返回:
#         dict: 融合后的结果
#     """
#     # 创建一个字典来存储每个文档的融合分数
#     fused_scores = {}
#
#     # 处理Chroma结果
#     if chroma_results and 'ids' in chroma_results and chroma_results['ids']:
#         # 获取Chroma的ID和距离
#         # Chroma返回的结果格式：{'ids': [['id1', 'id2']], 'distances': [[0.5, 0.8]], 'documents': [['doc1', 'doc2']]}
#         ids = chroma_results['ids'][0] if isinstance(chroma_results['ids'], list) and len(chroma_results['ids']) > 0 else chroma_results['ids']
#         distances = chroma_results['distances'][0] if 'distances' in chroma_results and isinstance(chroma_results['distances'], list) and len(chroma_results['distances']) > 0 else []
#
#         # 确保ids是列表形式
#         if not isinstance(ids, list):
#             ids = [ids]
#         if not isinstance(distances, list):
#             distances = [distances] if distances else []
#
#         # 距离转换为分数（距离越小分数越高）
#         for i, (doc_id, distance) in enumerate(zip(ids, distances)):
#             # 确保doc_id是字符串形式
#             doc_id = str(doc_id)
#
#             score = 1 / (distance + 1)  # 距离转分数
#             fused_scores[doc_id] = fused_scores.get(doc_id, 0) + 1 / (k + i + 1) * score
#
#     # 处理Elasticsearch结果
#     if es_results and 'ids' in es_results and es_results['ids']:
#         ids = es_results['ids']
#         scores = es_results['scores']
#
#         # 确保ids和scores是列表形式
#         if not isinstance(ids, list):
#             ids = [ids]
#         if not isinstance(scores, list):
#             scores = [scores]
#
#         for i, (doc_id, score) in enumerate(zip(ids, scores)):
#             # 确保doc_id是字符串形式
#             doc_id = str(doc_id)
#
#             # Elasticsearch的分数通常较高，进行归一化处理
#             normalized_score = score / 10 if score > 10 else score
#             fused_scores[doc_id] = fused_scores.get(doc_id, 0) + 1 / (k + i + 1) * normalized_score
#
#     # 按融合分数排序
#     sorted_results = sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)
#
#     # 构造返回结果
#     fused_results = {
#         'ids': [str(item[0]) for item in sorted_results],
#         'scores': [item[1] for item in sorted_results]
#     }
#
#     # 为融合结果添加文档内容
#     fused_results['documents'] = []
#
#     # 安全地提取Chroma文档
#     chroma_docs = {}
#     if chroma_results and 'ids' in chroma_results and 'documents' in chroma_results:
#         chroma_ids = chroma_results['ids'][0] if isinstance(chroma_results['ids'], list) and len(chroma_results['ids']) > 0 else chroma_results['ids']
#         chroma_documents = chroma_results['documents'][0] if isinstance(chroma_results['documents'], list) and len(chroma_results['documents']) > 0 else chroma_results['documents']
#
#         # 确保都是列表形式
#         if not isinstance(chroma_ids, list):
#             chroma_ids = [chroma_ids]
#         if not isinstance(chroma_documents, list):
#             chroma_documents = [chroma_documents]
#
#         if len(chroma_ids) == len(chroma_documents):
#             # 确保ID是字符串形式
#             chroma_ids = [str(id) for id in chroma_ids]
#             chroma_docs = dict(zip(chroma_ids, chroma_documents))
#
#     # 安全地提取ES文档
#     es_docs = {}
#     if es_results and 'ids' in es_results and 'documents' in es_results:
#         es_ids = es_results['ids']
#         es_documents = es_results['documents']
#
#         # 确保都是列表形式
#         if not isinstance(es_ids, list):
#             es_ids = [es_ids]
#         if not isinstance(es_documents, list):
#             es_documents = [es_documents]
#
#         if len(es_ids) == len(es_documents):
#             # 确保ID是字符串形式
#             es_ids = [str(id) for id in es_ids]
#             es_docs = dict(zip(es_ids, es_documents))
#
#     # 为每个融合的ID查找对应文档
#     for doc_id in fused_results['ids']:
#         doc_content = chroma_docs.get(doc_id) or es_docs.get(doc_id)
#         if doc_content:
#             fused_results['documents'].append(doc_content)
#
#     return fused_results
#
#
# def normalize_scores(scores):
#     """
#     对分数进行归一化处理
#
#     参数:
#         scores (list): 分数列表
#     返回:
#         list: 归一化后的分数列表
#     """
#     if not scores:
#         return []
#
#     min_score = min(scores)
#     max_score = max(scores)
#
#     # 如果所有分数相同，则返回相同值的列表
#     if min_score == max_score:
#         return [0.5 for _ in scores]
#
#     # 归一化公式: (score - min) / (max - min)
#     normalized_scores = [(score - min_score) / (max_score - min_score) for score in scores]
#     return normalized_scores


def hybrid_query(query_text, collection_name="employee_manual", es_index_name="health_advice", n_results=5):
    """
    混合检索：同时查询Chroma DB和Elasticsearch，并融合结果

    参数:
        query_text (str): 用户的问题文本
        collection_name (str): Chroma集合名称
        es_index_name (str): Elasticsearch索引名称
        n_results (int): 返回结果数量
    返回:
        dict: 融合后的检索结果
    """
    # 从Chroma DB检索
    chroma_results = query_chroma(query_text, collection_name, n_results)
    # print(chroma_results,"Chroma DB检索")

    # 从Elasticsearch检索
    # es_results = query_elasticsearch(query_text, es_index_name, n_results)
    # # print(es_results,"Elasticsearch检索")
    #
    # # 融合两个检索系统的分数
    # fused_results = reciprocal_rank_fusion(chroma_results, es_results)
    #
    # # 只返回前n_results个结果
    # final_results = {
    #     'ids': fused_results['ids'][:n_results],
    #     'documents': fused_results['documents'][:n_results],
    #     'scores': normalize_scores(fused_results['scores'][:n_results])
    # }
    #
    # return final_results
    return chroma_results


def get_llm_response(prompt):
    """
    调用阿里云百炼大模型生成回答

    参数:
        prompt (str): 发送给大模型的提示词
    返回:
        str: 大模型生成的回答文本
    """
    # 初始化OpenAI兼容客户端（连接阿里云百炼服务）
    client = OpenAI(
        api_key=os.getenv("DASHSCOPE_API_KEY"),  # 从环境变量获取API密钥
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",  # 百炼API地址
    )

    # 调用大模型生成对话响应
    completion = client.chat.completions.create(
        model="qwen-plus",  # 指定使用的大模型（阿里云通义千问plus版本）
        messages=[{"role": "user", "content": prompt}],  # 输入对话消息
    )
    # 返回生成的回答内容
    return completion.choices[0].message.content


def build_prompt_with_context(question, context_chunks):
    """
    根据用户问题和检索到的上下文构建提示词

    参数:
        question (str): 用户的问题
        context_chunks (list): 检索到的文本块列表
    返回:
        str: 构建好的提示词
    """
    # 构建上下文（将文本块格式化，便于大模型理解）
    context = "\n".join([f"- {chunk}" for chunk in context_chunks])
    
    # 构建提示词（明确大模型的角色、参考内容、问题和回答要求）
    prompt = f"""你是一个专业的旅行服务助手，请仔细阅读以下从知识库中检索到的所有相关内容，并准确回答用户问题。

    【知识库检索结果】：
    以下内容来自知识库中的多个文档，请综合使用这些信息来回答问题：

            {context}

            【用户问题】：
            {question}

            【回答要求】：
    1. **综合使用所有检索到的文档内容**：请仔细阅读上述所有参考内容，无论它们来自哪个文档（如12306.docx、铁路客运规章制度一览.docx等），都可以用来回答问题
    2. **严格基于知识库内容**：只使用上述检索到的内容回答，不要编造、推测或使用知识库外的信息
    3. **如果信息不足**：如果检索到的内容中确实没有相关信息，请明确说明"根据当前检索到的知识库内容，未找到相关信息"
    4. **回答格式**：
       - 回答应简洁明了，条理清晰
       - 如有多个要点，请分点列出（使用序号或项目符号）
       - 如果涉及具体条款或规定，可引用相关内容
       - 如果信息来自多个文档，可以综合说明
    5. **准确性**：确保回答准确，避免误导用户

    请基于上述所有检索到的内容，为用户提供准确、完整的回答。"""
    return prompt


def answer_user_question(question):
    """
    回答用户问题的主要函数（RAG流程）

    参数:
        question (str): 用户的问题
    返回:
        str: 最终的回答
    """
    # 1. 使用混合检索获取相关文档片段
    # results = hybrid_query(question)
    results = query_chroma(question)
    # 2. 从结果中提取文档内容
    if results and results['documents']:
        # 提取检索到的文本块并去重
        retrieved_chunks = list(set(results["documents"]))
    else:
        retrieved_chunks = []
    
    # 3. 如果没有找到相关内容，给出默认回复
    if not retrieved_chunks:
        return "抱歉，在知识库中未找到与您问题相关的信息。"
    
    # 4. 构建带有上下文的提示词
    prompt = build_prompt_with_context(question, retrieved_chunks)
    
    # 5. 调用大模型获得答案
    answer = get_llm_response(prompt)
    
    return answer


def evaluate_rag_response(question, ground_truth_answer, generated_answer, retrieved_chunks):
    """
    评估RAG系统的性能，计算准确率、召回率和F1值

    参数:
        question (str): 用户问题
        ground_truth_answer (str): 标准答案
        generated_answer (str): 模型生成的答案
        retrieved_chunks (list): 检索到的文本块

    返回:
        dict: 包含准确率、召回率和F1值的字典
    """
    # 构建评估提示词
    context = "\n".join([f"- {chunk}" for chunk in retrieved_chunks])
    evaluation_prompt = f"""
    请根据以下标准评估AI助手对问题的回答质量，并给出准确率、召回率和F1值的评分（每个指标满分1.0，保留2位小数）：
    
    评估标准：
    1. 准确性：回答的内容是否准确无误，是否与参考资料一致
    2. 相关性：回答是否与用户问题相关
    3. 完整性：回答是否充分覆盖了问题的所有方面
    
    参考资料：
    {context}
    
    用户问题：
    {question}
    
    标准答案：
    {ground_truth_answer}
    
    AI助手的回答：
    {generated_answer}
    
    请按照以下格式返回结果：
    准确率: [数值]
    召回率: [数值]
    F1值: [数值]
    
    只需给出这三个数值即可，不需要解释。
    """

    # 调用大模型进行评估
    try:
        evaluation_result = get_llm_response(evaluation_prompt)
        
        # 解析评估结果
        accuracy = 0.0
        recall = 0.0
        f1 = 0.0

        # 从响应中提取数值
        lines = evaluation_result.strip().split('\n')
        for line in lines:
            if '准确率:' in line:
                try:
                    accuracy = round(float(line.split(':')[1].strip()), 2)
                except:
                    pass
            elif '召回率:' in line:
                try:
                    recall = round(float(line.split(':')[1].strip()), 2)
                except:
                    pass
            elif 'F1值:' in line:
                try:
                    f1 = round(float(line.split(':')[1].strip()), 2)
                except:
                    pass

        return {
            "accuracy": accuracy,
            "recall": recall,
            "f1": f1
        }
    except Exception as e:
        print(f"评估过程中出现错误: {e}")
        return {
            "accuracy": 0.0,
            "recall": 0.0,
            "f1": 0.0
        }


def rag_with_evaluation(question, ground_truth_answer=None):
    """
    带评估的RAG流程

    参数:
        question (str): 用户问题
        ground_truth_answer (str): 标准答案（可选）

    返回:
        dict: 包含答案和评估结果的字典
    """
    # 1. 使用混合检索获取相关文档片段
    results = hybrid_query(question)
    # results = query_chroma(question)
    # 2. 从结果中提取文档内容
    # if results and results['documents']:
    #     # 提取检索到的文本块并去重
    #     retrieved_chunks = list(set(results["documents"]))
    # else:
    #     retrieved_chunks = []
    #
    # # 3. 如果没有找到相关内容，给出默认回复
    # if not retrieved_chunks:
    #     answer = "抱歉，在知识库中未找到与您问题相关的信息。"
    #     if ground_truth_answer:
    #         evaluation = evaluate_rag_response(question, ground_truth_answer, answer, [])
    #         return {"answer": answer, "evaluation": evaluation}
    #     else:
    #         return {"answer": answer, "evaluation": None}
    #
    # 4. 构建带有上下文的提示词
    prompt = build_prompt_with_context(question, results['documents'])
    
    # 5. 调用大模型获得答案
    answer = get_llm_response(prompt)
    
    # 6. 如果提供了标准答案，则进行评估
    if ground_truth_answer:
        evaluation = evaluate_rag_response(question, ground_truth_answer, answer, results['documents'])
        return {"answer": answer, "evaluation": evaluation}
    else:
        return {"answer": answer, "evaluation": None}


if __name__ == "__main__":
    # 测试问答功能mo
    user_question = "购票渠道。"
    ground_truth = """旅客可以通过多种方式购买铁路客运车票。常见的有前往火车站售票窗口，直接与售票员沟通，告知出行日期、目的地等信息，购买所需车票。也可以在火车站的自助售票机上操作，按照屏幕提示，选择车次、座位类型等进行购票。此外，各大正规售票网站和手机 APP 也是便捷的购票途径，比如 12306 官方网站及 APP，旅客只需在上面注册账号，登录后即可轻松查询车次、余票情况，并完成购票流程。"""  # 示例标准答案
    
    print(f"用户问题: {user_question}")
    print("-" * 50)
    
    result = rag_with_evaluation(user_question, ground_truth)
    print(f"AI助手回答:\n{result['answer']}")
    
    if result['evaluation']:
        print("\n" + "="*50)
        print("RAG系统评估结果:")
        print(f"准确率: {result['evaluation']['accuracy']}")
        print(f"召回率: {result['evaluation']['recall']}")
        print(f"F1值: {result['evaluation']['f1']}")